/* Copyright (C) 2010 Luca Piccinelli
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * File:   main.cpp
 * Author: lpiccinelli
 *
 * Created on 3 agosto 2010, 21.34
 */

#include <cstdlib>
#include <iostream>
#include <algorithm>

#include <boost/any.hpp>

#include <cv.h>
#include <highgui.h>

#include "prj_definitions.h"
#include "pipeline/adv_pipeline_data_structures.h"
#include "pipeline/pipeline_data_structures.h"
#include "pipeline/pipeline_io.h"
#include "generators/genericGenerator.h"
#include "generators/resource/CvImgDirectoryGenerator.h"
#include "generators/resource/D2DNBayesModelFileGenerator.h"
#include "generators/bridgeGenerator.h"
#include "generators/cameraInputGenerator.h"
#include "formatters/cvMat_formatters.h"
#include "formatters/ml_formatters.h"
#include "imgops/filters/GaussBlurOp.h"
#include "imgops/CvWhiteBalanceOp.h"
#include "imgops/skin/SkinTrainNBayesOp.h"
#include "imgops/skin/AdaptiveSkinRecNBayesOp.h"
#include "imgops/thresholding/CvHysteresisThreshOp.h"
#include "factory/factories.h"

using namespace std;
using namespace cv;
using namespace NAMESPACE;

class IsOpen{
public:
    CameraInputGenerator* g;

    bool operator()() const{
        if(waitKey(20) >= 0) g->close();
        return g->is_opened();
    }
};

class SwitchModel{
public:
    SwitchModel(int _switch_time) : switch_time(_switch_time){}

    int switch_time;

    LoopPipeStep<IsOpen>* loop;
    bool* update;

    NormalBayesModel *m1, *m2;

    bool operator()() const{
        int loop_num = loop->get_loops();
        *update = loop_num % switch_time;
        if(!(*update) && loop_num)
            m1->swap(*m2);

        return loop_num;
    }
};

int main(int argc, char** argv) {
    cout << "start running" << endl;
    
    string model_path;
    model_path = argv[1];

    // Statistics
    clock_t start_time, end_time;
    float framerate = 0, seconds = 0;

    FactoriesPool<Generator>* factoriesPool = FactoriesPool<Generator>::get_instance();

    // ****** Operations *******************************************************
    CvWhiteBalanceOp wb_op;
    SkinTrainNBayesOp skin_train;
    AdaptiveSkinRecNBayesOp skin_rec;
    CvHysteresisThreshOp threshold;
    // -------------------------------------------------------------------------

    // ****** Outputs **********************************************************
    // white balance
    GenericGenerator<Mat> outImgsGenerator(wb_op.getOutputKeysList().at(0));
    // skin train
    NormalBayesModel* online_model_train = new Discret2DNormalBayesModel(0, 0, 255, 255);
    HandledIOElement  online_model_train_io(online_model_train, skin_train.getOutputKeysList().at(0));
    // skin_recognition
    GenericGenerator<Mat> skin_out_img_generator(skin_rec.getOutputKeysList().at(0));
    // threshold
    GenericGenerator<Mat> thresh_out_img_generator(threshold.getOutputKeysList().at(0));
    // -------------------------------------------------------------------------


    // ****** White balance ****************************************************

    // white balance input
    ConfigurationTemplate cig_cfg_tmplt = factoriesPool->get_configuration_template("CameraInputGenerator");
    Configuration cig_cfg(cig_cfg_tmplt);
    cig_cfg.set_parameter(cig_cfg_tmplt.param_at(0), wb_op.getInputKeysList().at(0));
    cig_cfg.set_parameter(cig_cfg_tmplt.param_at(1), 320);
    cig_cfg.set_parameter(cig_cfg_tmplt.param_at(2), 240);
    cig_cfg.set_parameter(cig_cfg_tmplt.param_at(3), 0);
    Generator* imgsGenerator = factoriesPool->build("CameraInputGenerator", cig_cfg);
    // white balance criteria
    WB_METHOD wb_method = WB_WHITE_POINT;
    HandledIOElement wb_method_io(&wb_method, wb_op.getCriteriaKeysList().at(0));

    // ****** Output formatters ************************************************
    CvMatHgWndFormatter out_img_formatter(wb_op.getOutputKeysList().at(0), "white balance");
    // -------------------------------------------------------------------------

    // ****** White balance Pipestep *******************************************
    PipeStep wb_step;
    wb_step.set_computation(&wb_op);
    wb_step.add_input_generator(imgsGenerator);      // Input
    wb_step.add_criteria(&wb_method_io);             // Criteria
    wb_step.add_output_generator(&outImgsGenerator); // Output

    wb_step.add_output_formatter(&out_img_formatter);
    // -------------------------------------------------------------------------

    // *************************************************************************


    CVT_SPACE cvt_space = YCRCB;

    // ****** Skin recognition *************************************************
    // skin_recognition input
    BridgeGenerator skin_img_generator(outImgsGenerator, skin_rec.getInputKeysList().at(0));
    D2DNBayesModelFileGenerator modelGenerator(model_path, skin_rec.getInputKeysList().at(1));
    IOElement* model_off_io = modelGenerator.generate();
    NormalBayesModel* model_on = new Discret2DNormalBayesModel(0, 0, 255, 255);
    HandledIOElement model_on_io = HandledIOElement(model_on, skin_rec.getInputKeysList().at(2));
    // skin_recognition criteria
    NumericIOElement gain_io(0.8, skin_rec.getCriteriaKeysList().at(0));
    HandledIOElement cvt_space_io(&cvt_space, skin_rec.getCriteriaKeysList().at(1));
    int skin_model_class = 1;
    HandledIOElement skin_model_class_io(&skin_model_class, skin_rec.getCriteriaKeysList().at(2));

    // ****** Output formatters ************************************************
    CvMatHgWndFormatter skin_out_img_formatter(skin_rec.getOutputKeysList().at(0), "skin recognition");
    // -------------------------------------------------------------------------

    // ****** Skin recognition Pipestep ****************************************
    PipeStep skin_rec_step;
    skin_rec_step.set_computation(&skin_rec);
    skin_rec_step.add_input_generator(imgsGenerator)                   // Input
                 .add_input(model_off_io)
                 .add_input(&model_on_io);
    skin_rec_step.add_criteria(&gain_io)                               // Criteria
                 .add_criteria(&cvt_space_io)
                 .add_criteria(&skin_model_class_io);
    skin_rec_step.add_output_generator(&skin_out_img_generator);       // Output

    skin_rec_step.add_output_formatter(&skin_out_img_formatter);
    // -------------------------------------------------------------------------

    // *************************************************************************



    // ****** Thresholding *****************************************************
    // Input
    BridgeGenerator thresh_generator(skin_out_img_generator, threshold.getInputKeysList().at(0));
    //Criteria
    NumericIOElement thresh_t1_io(50.0,  threshold.getCriteriaKeysList().at(0));
    NumericIOElement thresh_t2_io(150.0, threshold.getCriteriaKeysList().at(1));
    NumericIOElement thresh_v1_io(0,   threshold.getCriteriaKeysList().at(2));
    NumericIOElement thresh_v2_io(255, threshold.getCriteriaKeysList().at(3));
    NumericIOElement thresh_block_size_io(3, threshold.getCriteriaKeysList().at(4));

    // ****** Output formatters ************************************************
    CvMatHgWndFormatter thresh_out_img_formatter(threshold.getOutputKeysList().at(0), "skin threshold");
    // -------------------------------------------------------------------------

    // ****** Pipestep *********************************************************
    PipeStep thresh_step;
    thresh_step.set_computation(&threshold);
    thresh_step.add_input_generator(&thresh_generator);  // Input
    thresh_step.add_criteria(&thresh_t1_io)
               .add_criteria(&thresh_t2_io)
               .add_criteria(&thresh_v1_io)
               .add_criteria(&thresh_v2_io)
               .add_criteria(&thresh_block_size_io);
    thresh_step.add_output_generator(&thresh_out_img_generator); // Output

    thresh_step.add_output_formatter(&thresh_out_img_formatter);
    // -------------------------------------------------------------------------

    // *************************************************************************



    // ****** Online Skin training *********************************************
    // input
    ConfigurationTemplate stig_cfg_tmplt = factoriesPool->get_configuration_template("BridgeGenerator");
    Configuration stig_cfg(stig_cfg_tmplt);
    stig_cfg.set_parameter(stig_cfg_tmplt.param_at(0), imgsGenerator);
    stig_cfg.set_parameter(stig_cfg_tmplt.param_at(1), skin_train.getInputKeysList().at(0));
    stig_cfg.set_parameter(stig_cfg_tmplt.param_at(2), 1);
    Generator* skin_train_img_generator = factoriesPool->build("BridgeGenerator", stig_cfg);

    //BridgeGenerator skin_train_img_generator(*imgsGenerator, skin_train.getInputKeysList().at(0), 1);
    BridgeGenerator skin_train_mask_generator(thresh_out_img_generator, skin_train.getInputKeysList().at(1));
    // criteria
    HandledIOElement train_cvt_space_io(&cvt_space, skin_train.getCriteriaKeysList().at(0));
    bool update_skin_model = true;
    HandledIOElement update_skin_model_io(&update_skin_model, skin_train.getCriteriaKeysList().at(1));

    BayesHgWndFormatter online_model_img_formatter(skin_train.getOutputKeysList().at(0));

    // Pipestep
    ConditionalPipeStep<SwitchModel> skin_train_step;
    skin_train_step.set_computation(&skin_train);
    skin_train_step.add_input_generator(skin_train_img_generator)   // Input
                   .add_input_generator(&skin_train_mask_generator);
    skin_train_step.add_criteria(&train_cvt_space_io)                // Criteria
                   .add_criteria(&update_skin_model_io);
    skin_train_step.add_output(&online_model_train_io);              // Output

    //skin_train_step.add_output_formatter(&online_model_img_formatter);
    // *************************************************************************

    // Step of pipeline looping
    LoopPipeStep<IsOpen> main_step;

    IsOpen is_open;
    is_open.g = dynamic_cast<CameraInputGenerator*>(imgsGenerator);

    SwitchModel switch_model(50);
    switch_model.loop = &main_step;
    switch_model.update = &update_skin_model;
    switch_model.m1 = model_on;
    switch_model.m2 = online_model_train;

    // Pipeline
    skin_train_step.set_condition(switch_model);

    Pipeline main_pipe;
    //main_pipe.add_pipe_step(&wb_step);
    main_pipe.add_pipe_step(&skin_rec_step);
    main_pipe.add_pipe_step(&thresh_step);
    main_pipe.add_pipe_step(&skin_train_step);

    main_step.set_condition(is_open);
    main_step.set_computation(&main_pipe);

    dynamic_cast<CameraInputGenerator*>(imgsGenerator)->open();
    start_time = clock();
    main_step.execute();
    end_time = clock();

    // ****** computing statistics *********************************************
    seconds = static_cast<float>(end_time - start_time) / static_cast<float>(CLOCKS_PER_SEC);
    framerate = static_cast<float>(main_step.get_loops()) / seconds;
    // -------------------------------------------------------------------------

    cout << framerate << " f/second" << endl;

    // ****** free resources ***************************************************
    delete imgsGenerator;
    delete skin_train_img_generator;
    delete online_model_train;
    delete model_on;
    // -------------------------------------------------------------------------

    return 0;
}